'''
@Time    : 2022/3/21 20:21
@Author  : Fu Junyu
@Site    : www.fujunyu.cn
@File    : datalodar.py
@Software: PyCharm
'''
import dgl
import numpy as np
import torch
from utils.decoder import occupationDict, occupationDictTest


class datalodar():
    def __init__(self, dataset):
        if dataset == 'ML':
            self.train_file = 'Ml-100k/u1.base'

        self.trainUser = []
        self.trainItem = []
        self.trainRating = []
        self.etypes = []

    def forward(self):
        with open(self.train_file) as f:
            for l in f.readlines():
                if len(l) > 0:
                    l = l.strip('\n').split('	')
                    items = int(l[1])-1  # dglGraph节点索引从0开始，但数据集中从1开始，所以将这里的索引减1
                    rating = float(str(l[2]))
                    uid = int(l[0])-1  # dglGraph节点索引从0开始，但数据集中从1开始，所以将这里的索引减1
                    time = int(l[3])

                    self.trainRating.append(rating)
                    self.trainUser.append(uid)
                    self.trainItem.append(items)

        self.etypes = set(self.trainRating)
        self.trainRating = np.array(self.trainRating)

        self.rels = {}

        for e in self.etypes:
            edges = np.argwhere(self.trainRating==e).flatten()
            self.rels[str(e)] = [[], []]
            for edge in edges:
                self.rels[str(e)][0].append(self.trainUser[edge])
                self.rels[str(e)][1].append(self.trainItem[edge])

        graph_data = {}

        for etype in self.rels:
            graph_data[('user', etype, 'movies')] = (torch.LongTensor(self.rels[etype][0]), torch.LongTensor(self.rels[etype][1]))
            graph_data[('movies', 're_%s' % etype, 'user')] = (torch.LongTensor(self.rels[etype][1]), torch.LongTensor(self.rels[etype][0]))

        # graph_data[('user', 'user_self', 'user')] = (torch.LongTensor(list(set(self.trainUser))), torch.LongTensor(list(set(self.trainUser))))
        # graph_data[('movies', 'movies_self', 'movies')] = (torch.LongTensor(list(set(self.trainItem))), torch.LongTensor(list(set(self.trainItem))))

        graph = dgl.heterograph(graph_data)

        return graph


class userFeat():
    def __init__(self, dataset):
        super(userFeat, self).__init__()
        if dataset == 'ML':
            self.train_file = 'Ml-100k/u.user'

        self.userID = []
        self.userAge = []
        self.userGender = []
        self.userOccupation = []
        self.userZip = []

    @property
    def getFeat(self):
        genderDict = {'M': 0.0, 'F': 1.0}

        occupationDict = self.occupationDict
        occupationDict = occupationDict.Dict

        with open(self.train_file) as f:
            for l in f.readlines():
                if len(l) > 0:
                    l = l.strip('\n').split('|')

                    uid = float(l[0])-1
                    uage = float(l[1])
                    ugender = genderDict[str(l[2])]
                    uoccupation = occupationDict[str(l[3])]
                    if not str(l[4]).isdigit():
                        uzip = 0
                    else:
                        uzip = float(l[4])

                    self.userID.append(uid)
                    self.userAge.append(uage)
                    self.userGender.append(ugender)
                    self.userOccupation.append(uoccupation)
                    self.userZip.append(uzip)
        feat = torch.tensor([self.userID, self.userAge, self.userGender, self.userOccupation, self.userZip])
        feat = feat.T
        return feat

    @property
    def occupationDict(self):
        return occupationDict()


class moviesFeat():
    def __init__(self, dataset):
        super(moviesFeat, self).__init__()

        if dataset == 'ML':
            self.train_file = 'Ml-100k/u.item'

    @property
    def getFeat(self):
        feat = []

        with open(self.train_file, "r", encoding='ISO-8859-1') as f:
            for l in f.readlines():
                if len(l) > 0:
                    l = l.strip('\n').split('|')
                    mfeat = [float(l[0])-1]
                    mtype = list(map(float, l[5:]))
                    mfeat.extend(mtype)
                    feat.append(mfeat)
                    # print(mfeat)
        feat = torch.tensor(feat)
        return feat



class trainData():
    def __init__(self, dataset):
        super(trainData, self).__init__()

        if dataset == 'ML':
            self.train_file = 'Ml-100k/u1.base'
        elif dataset =='ML_test':
            self.train_file = 'Ml-100k/u1.test'

    @property
    def getData(self):
        userID, movieID, labels = [], [], []

        with open(self.train_file, "r", encoding='ISO-8859-1') as f:
            for l in f.readlines():

                if len(l) > 0:
                    l = l.strip('\n').split('	')
                    uid = int(l[0])-1
                    mid = int(l[1])-1
                    rating = float(l[2])

                    userID.append(uid)
                    movieID.append(mid)
                    labels.append(rating)

        # return {'userID': torch.tensor(userID), 'moviesID': torch.tensor(movieID), 'labels': torch.tensor(labels)}
        if self.train_file == 'Ml-100k/u1.base':
            return self.batch(userID, movieID, labels, 3000)
        elif self.train_file == 'Ml-100k/u1.test':
            return self.batch(userID, movieID, labels, 2000)


    def batch(self, userID, movieID, labels, batchsize):  # userFeat为列表,平分后每份列表的的个数n
        for i in range(0, len(userID), batchsize):
            yield {'userID': torch.tensor(userID[i:i + batchsize]), 'moviesID': torch.tensor(movieID[i:i + batchsize]),
                   'labels': torch.tensor(labels[i:i + batchsize])}


class userFeatTest():
    def __init__(self, dataset):
        super(userFeatTest, self).__init__()
        if dataset == 'ML':
            self.train_file = 'Ml-100k/u.user'

    @property
    def getFeat(self):

        userID = []
        userAge = []
        userGender = []
        userOccupation = []
        userZip = []

        genderDict = {'M': 0, 'F': 1}
        occupationDict = occupationDictTest()
        occupationDict = occupationDict.Dict

        with open(self.train_file) as f:
            for l in f.readlines():
                if len(l) > 0:
                    l = l.strip('\n').split('|')

                    uid = int(l[0]) - 1
                    uage = int(l[1])
                    ugender = genderDict[str(l[2])]
                    uoccupation = occupationDict[str(l[3])]
                    if not str(l[4]).isdigit():
                        uzip = 0
                    else:
                        uzip = int(l[4])

                    userID.append(uid)
                    userAge.append(uage)
                    userGender.append(ugender)
                    userOccupation.append(uoccupation)
                    userZip.append(uzip)

            userID = torch.tensor(userID)
            userAge = torch.tensor(userAge)
            userGender = torch.tensor(userGender)
            userOccupation = torch.tensor(userOccupation)
            userZip = torch.tensor(userZip)

            userInput = {'uid': userID, 'age': userAge, 'gender': userGender, 'job': userOccupation, 'zip': userZip}
            return userInput

class moviesFeatTest():
    def __init__(self, dataset):
        super(moviesFeatTest, self).__init__()

        if dataset == 'ML':
            self.train_file = 'Ml-100k/u.item'

    @property
    def getFeat(self):

        movieID = []
        movieType = []

        with open(self.train_file, "r", encoding='ISO-8859-1') as f:
            for l in f.readlines():
                if len(l) > 0:
                    l = l.strip('\n').split('|')
                    mid = int(l[0])-1
                    mtype = list(map(int, l[5:]))

                    movieID.append(mid)
                    movieType.append(mtype)

        movieID = torch.tensor(movieID)
        movieType = torch.tensor(movieType)
        movieText = torch.load('Ml-100k/textTensor/movieText.pt')

        movieInput = {'mid': movieID, 'mtype': movieType, 'mtext': movieText}

        return movieInput



















